#ifndef RBTREE_HPP
#define RBTREE_HPP
#include <iostream>
#include <stdexcept>
#include <stack>
#include "config.hpp"
namespace myh {
    template<class T>
    class RBTree {
    private:
        enum Color { RED, BLACK };
        class RBN {
        public:
            T data;
            Color color;
            RBN *left, *right, *parent;
            RBN(T val) : data(val), color(RED), left(nullptr), right(nullptr), parent(nullptr) {}
        };
        RBN* root;
        void leftRotate(RBN* x) {
            if(!x||!x->right||!x->parent) return;
            RBN* y=x->right;
            x->right=y->left;
            if(y->left) y->left->parent=x;
            y->parent=x->parent;
            if(!x->parent) root=y;
            else if(x==x->parent->left) x->parent->left=y;
            else x->parent->right=y;
            y->left = x;
            x->parent = y;
        }
        void rightRotate(RBN* y) {
            if(!y||!y->left||!y->parent) return;
            RBN* x=y->left;
            y->left=x->right;
            if(x->right) x->right->parent=y;
            x->parent=y->parent;
            if(!y->parent) root=x;
            else if(y==y->parent->left) y->parent->left=x;
            else y->parent->right=x;
            x->right=y,y->parent=x;
        }
        void fixInsert(RBN* k) {
            if(!k) return;
            RBN* u;
            while(k!=root&&k->parent->color==RED) {
                if(k->parent == k->parent->parent->right) {
                    u = k->parent->parent->left;
                    if(u&& u->color==RED) {
                        u->color=BLACK;
                        k->parent->color=BLACK;
                        k->parent->parent->color=RED;
                        k=k->parent->parent;
                    }else{
                        if(k==k->parent->left) {
                            k=k->parent;
                            rightRotate(k);
                        }
                        k->parent->color=BLACK;
                        k->parent->parent->color=RED;
                        leftRotate(k->parent->parent);
                    }
                } else {
                    u=k->parent->parent->right;
                    if (u != nullptr && u->color==RED) {
                        u->color=BLACK,k->parent->color=BLACK;
                        k->parent->parent->color=RED;
                        k=k->parent->parent;
                    } else {
                        if(k==k->parent->right) {
                            k=k->parent;
                            leftRotate(k);
                        }
                        k->parent->color=BLACK,k->parent->parent->color=RED;
                        rightRotate(k->parent->parent);
                    }
                }
            }
            root->color = BLACK;
        }
        static void print_func(T data,bool flag){
            if(flag) std::cout<<data<<' ';
            else std::cout<<std::endl;
        }
        static bool cmp_func(T a,T b){
            return a<b;
        }
        typedef void(*func_ptr)(T,bool);
        typedef bool(*func_ptr2)(T,T);
        func_ptr2 cmp_ptr;
    public:
        RBTree(func_ptr2 cmp=cmp_func):root(nullptr),cmp_ptr(cmp){}
        template<class Container>
        RBTree(const Container& container,func_ptr2 cmp=cmp_func):root(nullptr){
            cmp_ptr=cmp;
            for(auto& i:container) insert(i);
        }
        void insert(T val) {
            RBN* node=new RBN(val);
            RBN* y=nullptr,*x=root;
            while(x!=nullptr) {
                y=x;
                if(cmp_ptr(node->data,x->data)) x = x->left;
                else if(cmp_ptr(x->data,node->data)) x = x->right;
                else{
                    delete node;
                    return;
                }
            }
            node->parent=y;
            if(!y) root=node;
            else if(cmp_ptr(node->data,y->data)) y->left=node;
            else y->right=node;
            fixInsert(node);
        }
        T* get_ptr(T a){
            if(!root) return nullptr;
            RBN* curr=root;
            while(curr){
                if(cmp_ptr(a,curr->data)) curr=curr->left;
                else if(cmp_ptr(curr->data,a)) curr=curr->right;
                else return &curr->data;
            }
            return nullptr;
        }
        void preorder(func_ptr func=print_func){
            if(root==nullptr) return;
            std::stack<RBN*> stack;
            RBN* current=root,*last_visited=nullptr;
            while(current!=nullptr||!stack.empty()) {
                if(current!=nullptr) {
                    stack.push(current);
                    current = current->left;
                } else {
                    RBN* peek_node = stack.top();
                    if(peek_node->right != nullptr && last_visited != peek_node->right)
                        current = peek_node->right;
                    else {
                        func(peek_node->data,true);
                        last_visited=peek_node;
                        stack.pop();
                    }
                }
            }
            func(0,false);
        }
        void inorder(func_ptr func=print_func) {
            std::stack<RBN*> stack;
            RBN* current = root;
            while(current||!stack.empty()) {
                while(current) {
                    stack.push(current);
                    current = current->left;
                }
                current=stack.top();
                stack.pop();
                func(current->data,true);
                current=current->right;
            }
            func(0,false);
        }
        void postorder(func_ptr func=print_func) {
            if(!root) return;
            std::stack<RBN*> stack;
            std::stack<T> output;
            stack.push(root);
            while(!stack.empty()) {
                RBN* current=stack.top();
                stack.pop();
                output.push(current->data);
                if(current->left) stack.push(current->left);
                if(current->right) stack.push(current->right);
            }
            while (!output.empty()) {
                func(output.top(),true);
                output.pop();
            }
            func(0,false);
        }
        ~RBTree() {
            RBN* current=root;
            while(current) {
                if(current->left) current = current->left;
                else if(current->right) current = current->right;
                else {
                    RBN* parent=current->parent;
                    if(parent) {
                        if(current==parent->left) parent->left=nullptr;
                        else parent->right=nullptr;
                    }
                    delete current;
                    current=parent;
                }
            }
            root=nullptr;
        }
    };
}
#endif
